/******************************************************************************
 * Copyright 2022 The Untouch Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/
#include "modules/calibration/lidar_imu_calibrator/include/record_data.h"

#include <QDir>
#include <QString>

#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>

#include "modules/localization/msf/common/io/pcl_point_types.h"
#include "modules/localization/msf/common/util/time_conversion.h"

RecordData::RecordData()
    : thread_pool_(new apollo::localization::msf::ThreadPool) {
  apollo::cyber::Init("lidar_imu_calibrator");
}

RecordData::~RecordData() {}

bool RecordData::Start() {
  thread_pool_->start(2);
  if (!Init()) {
    return false;
  }
  return true;
}

void RecordData::Stop() {
  CloseChannel();
  thread_pool_->stop();
  if (odometry_loc_stream_.is_open()) {
    odometry_loc_stream_.close();
  }
  WritePcdTimestapFile();
}

void RecordData::SetOdometryTopic(const std::string& topic) {
  odometry_topic_ = topic;
}

void RecordData::SetLidarTopics(const std::vector<std::string>& topics) {
  lidar_topics_ = topics;
}

void RecordData::SetCacheFolderPath(const std::string& path) {
  cache_folder_path_ = path;
}

bool RecordData::Init() {
  if (odometry_topic_.empty() || lidar_topics_.empty()) {
    AERROR << "image topic or lidar topic is empty!";
    return false;
  }

  if (cache_folder_path_.empty()) {
    AERROR << "cache folder path is empty!";
    return false;
  }

  QDir dir;
  if (!dir.exists(QString::fromStdString(cache_folder_path_))) {
    if (!dir.mkpath(QString::fromStdString(cache_folder_path_))) {
      AERROR << "failed to create cache folder!";
      return false;
    }
  }

  CloseChannel();
  node_ = apollo::cyber::CreateNode("LidarImuCalibrator");
  if (!node_) {
    AERROR << "failed to create node!";
    return false;
  }

  std::string odometry_file_path =
      cache_folder_path_ + "/odometry_localization.csv";
  odometry_loc_stream_.open(odometry_file_path,
                            std::ofstream::out | std::ofstream::trunc);
  if (odometry_loc_stream_.is_open()) {
    odometry_loc_stream_
        << "timestamp,x,y,z,qx,qy,qz,qw,std_dev_x,std_dev_y,std_dev_z\n";
  } else {
    AERROR << "failed to open file[" << odometry_file_path << "]!";
    return false;
  }

  odometry_reader_ = node_->CreateReader<apollo::localization::Gps>(
      odometry_topic_,
      std::bind(&RecordData::OdometryCallback, this, std::placeholders::_1));

  for (auto& lidar_topic : lidar_topics_) {
    QString save_path =
        QString::fromStdString(cache_folder_path_) + "/" +
        QString::fromStdString(lidar_topic).replace(QChar('/'), QChar('_'));
    dir.setPath(save_path);
    if (dir.exists()) {
      dir.removeRecursively();
    }

    if (!dir.mkdir(save_path)) {
      AERROR << "failed to create folder[" << save_path.toStdString() << "]!";
      return false;
    }

    auto lidar_reader = node_->CreateReader<apollo::drivers::PointCloud>(
        lidar_topic, std::bind(&RecordData::LidarCallback, this,
                               std::placeholders::_1, save_path.toStdString()));
    if (!lidar_reader) {
      AERROR << "failed to create lidar[" << lidar_topic << "] reader!";
      return false;
    }
    lidar_readers_.push_back(lidar_reader);
  }
  return true;
}

void RecordData::CloseChannel() {
  if (odometry_reader_) {
    odometry_reader_.reset();
  }

  for (auto& lidar_reader : lidar_readers_) {
    if (lidar_reader) {
      lidar_reader.reset();
    }
  }
  lidar_readers_.clear();

  if (node_) {
    node_.reset();
  }
}

void RecordData::OdometryCallback(
    const std::shared_ptr<apollo::localization::Gps>& message) {
  WritePoseFile(message);
}

void RecordData::LidarCallback(
    const std::shared_ptr<apollo::drivers::PointCloud>& message,
    const std::string& save_path) {
  thread_pool_->commit(
      std::bind(&RecordData::WritePcdFile, this, message, save_path));
}

void RecordData::WritePoseFile(
    const std::shared_ptr<apollo::localization::Gps>& message) {
  odometry_loc_stream_ << std::fixed;
  odometry_loc_stream_ << message->header().timestamp_sec();
  odometry_loc_stream_ << ",";
  odometry_loc_stream_ << message->localization().position().x();
  odometry_loc_stream_ << ",";
  odometry_loc_stream_ << message->localization().position().y();
  odometry_loc_stream_ << ",";
  odometry_loc_stream_ << message->localization().position().z();
  odometry_loc_stream_ << ",";
  odometry_loc_stream_ << message->localization().orientation().qx();
  odometry_loc_stream_ << ",";
  odometry_loc_stream_ << message->localization().orientation().qy();
  odometry_loc_stream_ << ",";
  odometry_loc_stream_ << message->localization().orientation().qz();
  odometry_loc_stream_ << ",";
  odometry_loc_stream_ << message->localization().orientation().qw();
  odometry_loc_stream_ << ",0,0,0\n";  // std_x = 0;std_y = 0; double std_z = 0;
}

void RecordData::WritePcdFile(
    const std::shared_ptr<apollo::drivers::PointCloud>& message,
    const std::string& save_path) {
  pcl::PointCloud<apollo::localization::msf::velodyne::PointXYZIT> cloud;
  cloud.width = message->width();
  cloud.height = message->height();
  cloud.is_dense = false;
  cloud.points.resize(cloud.width * cloud.height);

  if (cloud.width == 0 || cloud.height == 0) {
    cloud.width = 1;
    cloud.height = message->point_size();
    cloud.points.resize(message->point_size());
  }

  for (size_t i = 0; i < static_cast<size_t>(cloud.points.size()); ++i) {
    cloud.points[i].x = message->point(i).x();
    cloud.points[i].y = message->point(i).y();
    cloud.points[i].z = message->point(i).z();
    cloud.points[i].intensity =
        static_cast<unsigned char>(message->point(i).intensity());
    cloud.points[i].timestamp = message->point(i).timestamp() * 1e-9;
  }
  std::stringstream ss;
  ss << save_path;
  ss << "/";
  ss << apollo::localization::msf::TimeToString(message->measurement_time());
  ss << ".pcd";
  pcl::io::savePCDFileBinaryCompressed(ss.str(), cloud);
  AddPcdTimestamps(save_path, message->measurement_time());
}

void RecordData::AddPcdTimestamps(const std::string& save_path,
                                  double timestamp) {
  std::lock_guard<std::mutex> locker(pcd_timestamps_mutex_);
  pcd_timestamps_[save_path].push_back(timestamp);
}

void RecordData::WritePcdTimestapFile() {
  auto map_iter = pcd_timestamps_.begin();
  while (map_iter != pcd_timestamps_.end()) {
    auto iter = map_iter->second.begin();
    std::ofstream file(map_iter->first + "/pcd_timestamp.csv",
                       std::ofstream::out | std::ofstream::trunc);
    file << "timestamp\n";
    std::sort(map_iter->second.begin(), map_iter->second.end());
    while (iter != map_iter->second.end()) {
      file << apollo::localization::msf::TimeToString(*iter);
      file << "\n";
      ++iter;
    }
    file.close();
    ++map_iter;
  }
  pcd_timestamps_.clear();
}
